Shader "Kun URP/Post Process/Volume Cloud"
{
    Properties
    {
        _MainTex("MainTex",2D) = "white"{}
		_HeightDensity("HeightDensity", 2D) = "white"{}
		_BaseTex("BaseTex", 3D) = "white" {}
		_BaseTile("BaseTile", float) = 3.0
		_DetailTex("Detail", 3D) = "white" {}
		_DetailTile("DetailTile", float) = 10.0
		_DetailStrength("DetailStrength", float) = 0.2
		_CurlNoise("CurlNoise", 2D) = "white"{}
		_CurlTile("CurlTile", float) = .2
		_CurlStrength("CurlStrength", float) = 1
		_CloudTopOffset("TopOffset",float) = 100
		_CloudSize("CloudSize", float) = 50000

		_CloudStartHeight("CloudStartHeight", float) = 2000
		_CloudEndHeight("CloudEndHeight", float) = 8000

		_CloudOverallDensity("CloudOverallDensity",float) = .1
		_CloudCoverageModifier("CloudCoverageModifier", float) = 1.0
		_CloudTypeModifier("CloudTypeModifier", float) = 1.0

		_WeatherTex("WeatherTex", 2D) = "white" {}
		_WeatherTexSize("WeatherTexSize", float) = 50000
		_WindDirection("WindDirection",Vector) = (1,1,0,0)
		_SilverIntensity("SilverIntensity",float) = .8
		_ScatteringCoefficient("ScatteringCoefficient",float) = .04
		_ExtinctionCoefficient("ExtinctionCoefficient",float) = .04
		_MultiScatteringA("MultiScatteringA",float) = 0.5
		_MultiScatteringB("MultiScatteringB",float) = 0.5
		_MultiScatteringC("MultiScatteringC",float) = 0.5
		_SilverSpread("SilverSpread",float) = .75
		_RaymarchOffset("RaymarchOffset", float) = 0.0
		_AmbientColor("AmbientColor", Color) = (1,1,1,1)
		_AtmosphereColor("AtmosphereColor" , Color) = (1,1,1,1)
		_AtmosphereColorSaturateDistance("AtmosphereColorSaturateDistance", float) = 80000
        
    }
    SubShader
    {
        Tags{
            "RenderPipeline"="UniversalRenderPipeline"
        }
        Cull Off ZWrite Off ZTest Always
        HLSLINCLUDE

        #include "Packages/com.unity.render-pipelines.universal/ShaderLibrary/Core.hlsl"
        #include "../../ShaderLibrary/PostProcess/CloudHelper.hlsl"

        // TEXTURE2D( _MainTex);
        // SAMPLER(sampler_MainTex);
        // TEXTURE2D( _CameraDepthTexture);
        // SAMPLER(sampler_CameraDepthTexture);
        
        ENDHLSL

        // pass 0, Render a undersampled buffer.
        pass
        {
            HLSLPROGRAM
            #pragma vertex vert
            #pragma fragment frag

            #define MIN_SAMPLE_COUNT 32
			#define MAX_SAMPLE_COUNT 32

            float _RaymarchOffset;	//raymarch offset by halton sequence, [0,1]
            float2 _TexelSize;	//Texelsize used to decide offset by bayer matrix.
            float4 _ProjectionExtents;
            
            struct Attributes
            {
                float4 positionOS   : POSITION;
                float2 uv           : TEXCOORD;
            };
            struct Varyings
            {
                float4 positionCS   : SV_POSITION;
                float4 screenPos    : TEXCOORD0;
                float2 vsray : TEXCOORD1;
            };

            Varyings vert(Attributes input)
            {
                Varyings output;
                output.positionCS = TransformObjectToHClip(input.positionOS.xyz);
                output.screenPos = ComputeScreenPos(output.positionCS);
                output.vsray = (2.0 * input.uv - 1.0) * _ProjectionExtents.xy + _ProjectionExtents.zw;;
                return output;
            }

            half4 frag(Varyings input) : SV_TARGET
            {
                float3 vspos = float3(input.vsray, 1.0);
				float4 worldPos = mul(unity_CameraToWorld, float4(vspos,1.0));
				worldPos /= worldPos.w;
                // return worldPos;
				
				// float sceneDepth = Linear01Depth(_DownsampledDepth.Sample(sampler_DownsampledDepth, (i.screenPos / i.screenPos.w).xy));
				//TODO: sceneDepth here is distance in camera z-axis, but should be radial distance.
				//For now it causes bug that when allow cloud front object is enabled, the cloud is different when rotating camera.
				float raymarchEnd = 1e8;	
				float3 viewDir = normalize(worldPos.xyz - _WorldSpaceCameraPos);
				int sample_count = lerp(MAX_SAMPLE_COUNT, MIN_SAMPLE_COUNT, abs(viewDir.y));	//dir.y ==0 means horizontal, use maximum sample count

				float2 screenPos = input.screenPos.xy / input.screenPos.w;
				int2 texelID = int2(fmod(screenPos/ _TexelSize , 3.0));	//Calculate a texel id to index bayer matrix.
										
				float bayerOffset = (bayerOffsets[texelID.x][texelID.y]) / 9.0f;	//bayeroffset between[0,1)
				float offset = -fmod(_RaymarchOffset + bayerOffset, 1.0f);			//final offset combined. The value will be multiplied by sample step in GetDensity.

				float intensity, distance;
				
				float density = GetDensity(worldPos, viewDir, raymarchEnd, sample_count, offset, /*out*/intensity, /*out*/distance);
				// return float4(intensity, intensity, intensity, density);
                return float4(intensity, distance, 1.0f, density);
                // return float4(1, 1, 1, 1);
            }

            
            ENDHLSL
        }

        // pass 1, blend undersampled image with history buffer to new buffer.
        pass
        {
            HLSLPROGRAM
            #pragma vertex vert
            #pragma fragment frag

            sampler2D _MainTex;						//history buffer.
			float4 _MainTex_TexelSize;
			sampler2D _UndersampleCloudTex;			//current undersampled tex.
			float4 _UndersampleCloudTex_TexelSize;

            float4x4 _PrevVP;	//View projection matrix of last frame. Used to temporal reprojection.
			float4 _ProjectionExtents;
			float2 _TexelSize;
            
            struct Attributes
            {
                float4 positionOS   : POSITION;
                float2 uv           : TEXCOORD;
            };
            struct Varyings
            {
                float4 positionCS   : SV_POSITION;
                float2 uv           : TEXCOORD0;
                float4 screenPos    : TEXCOORD1;
                float2 vsray        : TEXCOORD2;
            };

            Varyings vert(Attributes input)
            {
                Varyings output;
                output.positionCS = TransformObjectToHClip(input.positionOS.xyz);
                output.uv = input.uv;
                output.screenPos = ComputeScreenPos(output.positionCS);
                output.vsray = (2.0 * input.uv - 1.0) * _ProjectionExtents.xy + _ProjectionExtents.zw;
                return output;
            }

            //Get uv of wspos in history buffer.
			float2 PrevUV(float4 wspos, out half outOfBound) {
				float4 prevUV = mul(_PrevVP, wspos);
				prevUV.xy = 0.5 * (prevUV.xy / prevUV.w) + 0.5;
				half oobmax = max(0.0 - prevUV.x, 0.0 - prevUV.y);
				half oobmin = max(prevUV.x - 1.0, prevUV.y - 1.0);
				outOfBound = step(0, max(oobmin, oobmax));
				return prevUV;
			}

			//Code from https://zhuanlan.zhihu.com/p/64993622. Do AABB clip in TAA(clip to center).
			float4 ClipAABB(float4 aabbMin, float4 aabbMax, float4 prevSample)
			{
				// note: only clips towards aabb center (but fast!)
				float4 p_clip = 0.5 * (aabbMax + aabbMin);
				float4 e_clip = 0.5 * (aabbMax - aabbMin);
				float4 v_clip = prevSample - p_clip;
				float4 v_unit = v_clip / e_clip;
				float4 a_unit = abs(v_unit);
				float ma_unit = max(max(a_unit.x, max(a_unit.y, a_unit.z)), a_unit.w);

				if (ma_unit > 1.0)
					return p_clip + v_clip / ma_unit;
				else
					return prevSample;// point inside aabb
			}

            half4 frag(Varyings input) : SV_TARGET
            {
                float3 vspos = float3(input.vsray, 1.0);
				float4 worldPos = mul(unity_CameraToWorld, float4(vspos, 1.0f));
				worldPos /= worldPos.w;
				float4 raymarchResult = tex2D(_UndersampleCloudTex, input.uv);
				float distance = raymarchResult.y;		
				float intensity = raymarchResult.x;
				half outOfBound;
				float2 prevUV = PrevUV(mul(unity_CameraToWorld, float4(normalize(vspos) * distance, 1.0)), outOfBound);	//find uv in history buffer.
					
				{	//Do temporal reprojection and clip things.
					float4 prevSample = tex2D(_MainTex, prevUV);
					float2 xoffset = float2(_UndersampleCloudTex_TexelSize.x, 0.0f);
					float2 yoffset = float2(0.0f, _UndersampleCloudTex_TexelSize.y);

					float4 m1 = 0.0f, m2 = 0.0f;

					float validSampleCount = 9.0f;
					//The loop below calculates mean and variance used to calculate AABB.
					[unroll]
					for (int x = -1; x <= 1; x ++) {
						[unroll]
						for (int y = -1; y <= 1; y ++ ) {
							float4 val;
							if (x == 0 && y == 0) {
								val = raymarchResult;
								m1 += val;
								m2 += val * val;
							}
							else {
								float2 uv = input.uv + float2(x * _UndersampleCloudTex_TexelSize.x, y * _UndersampleCloudTex_TexelSize.y);
								val = tex2Dlod(_UndersampleCloudTex, float4(uv, 0.0, 0.0));
								m1 += val;
								m2 += val * val;
							}
						}
					}
					//Code from https://zhuanlan.zhihu.com/p/64993622.
					float gamma = 1.0f;
					float4 mu = m1 / validSampleCount;
					float4 sigma = sqrt(abs(m2 / validSampleCount - mu * mu));
					float4 minc = mu - gamma * sigma;
					float4 maxc = mu + gamma * sigma;
					prevSample = ClipAABB(minc, maxc, prevSample);	

					//Blend.
					raymarchResult = lerp(prevSample, raymarchResult, max(0.05f, outOfBound));
				}
				return 	raymarchResult;
            }

            
            ENDHLSL
        }

        // pass 2, Calculate lighting, blend final cloud image with final image.
        pass
        {
            HLSLPROGRAM
            #pragma vertex vert
            #pragma fragment frag

            sampler2D _MainTex;	//Final image without cloud.
			float4 _MainTex_TexelSize;	//Final image without cloud.
			sampler2D _CloudTex;	//The full resolution cloud tex we generated.
            sampler2D _CameraDepthTexture;
			float4 _CloudTex_TexelSize;
            float4 _ProjectionExtents;
            float3 _AtmosphereColor;
            half3 _AmbientColor;
            float _AtmosphereColorSaturateDistance;
            
            struct Attributes
            {
                float4 positionOS   : POSITION;
                float2 uv           : TEXCOORD;
            };
            struct Varyings
            {
                float4 positionCS   : SV_POSITION;
                float2 uv           : TEXCOORD0;
                float4 screenPos    : TEXCOORD1;
                float2 vsray        : TEXCOORD2;
            };

            Varyings vert(Attributes input)
            {
                Varyings output;
                output.positionCS = TransformObjectToHClip(input.positionOS.xyz);
                output.uv = input.uv;
                output.screenPos = ComputeScreenPos(output.positionCS);
                output.vsray = (2.0 * input.uv - 1.0) * _ProjectionExtents.xy + _ProjectionExtents.zw;
                return output;
            }

            half4 frag(Varyings input) : SV_TARGET
            {
                float3 vspos = float3(input.vsray, 1.0);
				float4 worldPos = mul(unity_CameraToWorld,float4(vspos,1.0));
				float3 viewDir = normalize(worldPos.xyz - _WorldSpaceCameraPos);

				half4 mcol = tex2D(_MainTex,input.uv);
                float4 currSample = tex2D(_CloudTex, input.uv);
                float depth = currSample.g;
				
                Light mainLight = GetMainLight();

				float3 sunColor;
                sunColor = mainLight.color;
                float4 result;
				result.rgb = currSample.r * sunColor + currSample.b *_AmbientColor * currSample.a;
				result.a = currSample.a;
                float atmosphericBlendFactor = exp(-depth / _AtmosphereColorSaturateDistance);
				result.rgb = lerp(_AtmosphereColor, result.rgb, saturate(atmosphericBlendFactor));
                float sceneDepth = Linear01Depth(tex2D(_CameraDepthTexture, (input.screenPos / input.screenPos.w).xy), _ZBufferParams);
				if (sceneDepth == 1.0f) {
					return half4(mcol.rgb * (1 - result.a) + result.rgb * result.a, 1);
				}
				else {
					return mcol;
				}
            }

            
            ENDHLSL
        }
    }
}
